import os
import re
import numpy as np
import cv2 as cv
from PyCmpltrtok.common import *
from PyCmpltrtok.common_np import *
import matplotlib.pyplot as plt
import PyCmpltrtok.data.cifar10.load_cifar10 as cifar10
from onnx_helper import ONNXClassifierWrapper


if '__main__' == __name__:

    def _main():
        trt_path = '/home/asuspei/PycharmProjects/AsusCondaP37Torch1101Cuda111/python_ai/category/onnx/torch2onnx/_save/vgg16_torch2onnx.py/v1.0/2022_06_17_12_54_40_476426-64.trt'

        sep('Load images')
        img_dir = '/home/asuspei/large_data/DL1/_many_files/cifar-10_pngs/cifar/test'
        M = 64
        img_paths = os.listdir(img_dir)[:M]
        input_arr = np.zeros((M, 3, 32, 32), dtype=np.uint8)
        labels = []
        regexp = re.compile(r'_([^_]+)$')
        for i, xname in enumerate(img_paths):
            print(i, xname)
            base = os.path.splitext(xname)[0]
            matcher = regexp.search(base)
            xpath = os.path.join(img_dir, xname)
            if matcher is None:
                raise Exception(f'{xpath} cannot tell the label!')
            label = matcher.group(1)
            labels.append(label)
            img = cv.imread(xpath, cv.IMREAD_COLOR)[:, :, ::-1]
            img = img.transpose(2, 0, 1)
            input_arr[i] = img
        input_arr = uint8_to_flt_by_lut(input_arr)

        sep('Load model by TRT onnx helper')
        print(trt_path)
        trt_model = ONNXClassifierWrapper(trt_path, [M, 10], target_dtype=np.float32)

        sep('Run it')
        output_arr = trt_model.predict(input_arr)
        check_np(output_arr, 'output_arr')
        output_arr = output_arr.argmax(axis=1).tolist()
        print(output_arr)
        label_names = cifar10.load(only_meta=True)
        print(label_names)
        print([(i, name) for i, name in enumerate(label_names)])
        output = [label_names[idx] for idx in output_arr]
        print(output)

        sep('Plot')
        plt.figure(figsize=(14, 6))
        plt.subplots_adjust(wspace=0.6)
        spr = 6
        spc = 12
        spn = 0
        m, r = spr * spc, 0
        for i in range(m):
            if i > M - 1:
                break
            spn += 1
            plt.subplot(spr, spc, spn)
            plt.axis('off')
            gt = labels[i]
            h = output[i]
            is_right = gt == h
            if is_right:
                r += 1
            plt.title(h if is_right else f'{h}(gt: {gt})', color='black' if is_right else 'red')
            img = input_arr[i].transpose(1, 2, 0)
            plt.imshow(img)
        print(f'Accuracy of demo: {r/M:.4f}')
        print('Check the plotting window and close it to continue ...')
        plt.show()

    _main()
